package cn.virens.common.datasource;

import cn.virens.common.exception.APIException;
import org.springframework.jdbc.datasource.AbstractDataSource;

import javax.sql.DataSource;
import java.sql.Connection;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.atomic.AtomicInteger;

public class DynamicDataSource extends AbstractDataSource {
    private final AtomicInteger counter = new AtomicInteger(0);

    private List<DataSource> slaveDataSource = new ArrayList<>();

    private DataSource masterDataSource;

    @Override
    public Connection getConnection(String username, String password) throws SQLException {
        return determineTargetDataSource().getConnection(username, password);
    }

    @Override
    public Connection getConnection() throws SQLException {
        return determineTargetDataSource().getConnection();
    }

    /**
     * 设置主数据库实例
     *
     * @param dataSource 数据库实例
     */
    public void setMasterDataSource(DataSource dataSource) {
        this.masterDataSource = dataSource;
    }

    /**
     * 添加从数据库实例
     *
     * @param dataSource 数据库实例
     */
    public void addSlaveDataSource(DataSource dataSource) {
        this.slaveDataSource.add(dataSource);
    }

    /**
     * 获取当前有效的数据库连接池实例
     *
     * @return 数据库连接池实例
     * @throws APIException 无
     */
    protected DataSource determineTargetDataSource() throws APIException {
        DataSourceType type = DataSourceUtils.get0();
        if (DataSourceType.SLAVE.equals(type)) {
            return getSlaveDataSource();
        }

        return masterDataSource;
    }

    /**
     * 轮询获取从数据库连接池实例
     *
     * @return 从数据库连接池实例
     * @throws APIException 无
     */
    protected DataSource getSlaveDataSource() throws APIException {
        int couid = Math.abs(counter.incrementAndGet());
        int size = Math.abs(slaveDataSource.size());

        return slaveDataSource.get(couid % size);
    }

}
